// Remote Desktop System - remote controlling of multiple PC's
// Copyright (C) 2000-2012 GravyLabs LLC

// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; version 2 of the License.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

// TCPSocket.cpp : implementation file
// Implementation of AsyncSocket derived class.
// This class only works with versions of windows that recognize the TransmitPackets API function call
//

#include "stdafx.h"
#include "TCPSocket.h"
#include "Packet.h"

CTCPSocket::CTCPSocket() : baseclass(),
	m_bListening(FALSE),m_hWndParent(NULL),TransmitPackets(NULL)
{
}

CTCPSocket::~CTCPSocket()
{
}

// Initialize the transmit packet function pointer
BOOL CTCPSocket::InitTP()
{
	// Create the transmit packet function pointer
	DWORD len;
	static const GUID guid = WSAID_TRANSMITPACKETS;
	if (WSAIoctl(m_hSocket,SIO_GET_EXTENSION_FUNCTION_POINTER,(void*)&guid,sizeof(guid),&TransmitPackets,sizeof(TransmitPackets),&len,NULL,NULL))
		return FALSE;
	return TRUE;
}

// Create a new socket
BOOL CTCPSocket::Create(int nSocketPort)
{
	BOOL bCreate = FALSE;
	if (nSocketPort == 0)
	{
		// Create the outbound connection
		bCreate = baseclass::Create();

		// Receive notifications for connections, reading data, and disconnections
		AsyncSelect(FD_CONNECT | FD_READ | FD_CLOSE);
	}
	else
	{
		// Create the listener socket
		bCreate = baseclass::Create(nSocketPort);

		// Receive notifications of a waiting connection, reading data, and disconnections
		AsyncSelect(FD_ACCEPT | FD_READ | FD_CLOSE);

		// Start listening
		if (bCreate)
			m_bListening = baseclass::Listen();
	}

	// Create the transmit packet function pointer
	if (!InitTP())
		bCreate = FALSE;

	return bCreate;
}

// Set the parent window handle for events to be posted back to the parent with WM_APP messags
void CTCPSocket::SetParent(HWND hWndParent)
{
	ASSERT(hWndParent != NULL);
	m_hWndParent = hWndParent;
}

// Make a connection using a outbound socket
BOOL CTCPSocket::Connect(LPCWSTR lpszHostAddress,UINT nHostPort)
{
	if (!m_bListening)
		return baseclass::Connect(lpszHostAddress,nHostPort);
	return FALSE;
}

// Shutdown the socket
BOOL CTCPSocket::ShutDown(int nHow)
{
	return baseclass::ShutDown(nHow);
}

// Notification of a pending inbound connection
void CTCPSocket::OnAccept(int nErrorCode)
{
	ASSERT(m_bListening == TRUE);
	baseclass::OnAccept(nErrorCode);

	// Send the custom WM_ACCEPTCONN message
	if (m_hWndParent != NULL)
	{
		// Send and wait for processing
		SendMessage(m_hWndParent,WM_ACCEPTCONN,(WPARAM)this,0);
	}
}

// Notification that the connection has been successfully made
void CTCPSocket::OnConnect(int nErrorCode)
{
	baseclass::OnConnect(nErrorCode);

	// Send the custom WM_MAKECONN message
	if (m_hWndParent != NULL)
	{
		// Send and wait for processing
		SendMessage(m_hWndParent,WM_MAKECONN,(WPARAM)this,(LPARAM)nErrorCode);
	}
}

// Notification of readiness to receive data
void CTCPSocket::OnReceive(int nErrorCode)
{
	ASSERT(m_bListening == FALSE);
	baseclass::OnReceive(nErrorCode);

	// Send the custom WM_RECEIVEDATA message
	if (m_hWndParent != NULL)
	{
		// Send and wait for processing
		SendMessage(m_hWndParent,WM_RECEIVEDATA,(WPARAM)this,0);
	}
}

// Notification that the connection has been closed
void CTCPSocket::OnClose(int nErrorCode)
{
	baseclass::OnClose(nErrorCode);

	// Shutdown and close the socket
	ShutDown();
	Close();

	// Send the custom WM_CLOSECONN message
	if (m_hWndParent != NULL)
	{
		// Send and wait for processing
		SendMessage(m_hWndParent,WM_CLOSECONN,(WPARAM)this,0);
	}
}

// Helper for sorting, since order really doesn't matter, use the memory address
bool CTCPSocket::operator < (const CTCPSocket & rhs)
{
	return (DWORD_PTR)this < (DWORD_PTR)(&rhs);
}

// Helper to send a packet of data
CTCPSocket & CTCPSocket::operator << (CPacket & rhs)
{
	// Enable blocking mode and prevent events
	DWORD dwArgument = 0;
	AsyncSelect(dwArgument);
	IOCtl(FIONBIO,&dwArgument);

	// Send the packet
	rhs.SerializePacket(this,true);

	// Disable blocking mode and reinstate events
	dwArgument = FD_CONNECT | FD_READ | FD_CLOSE;
	IOCtl(FIONBIO,&dwArgument);
	AsyncSelect(dwArgument);

	// Return a reference to "this" class for chaining
	return *this;
}

// Helper to receive a packet
CTCPSocket & CTCPSocket::operator >> (CPacket & rhs)
{
	// Enable blocking mode and prevent events
	DWORD dwArgument = 0;
	AsyncSelect(dwArgument);
	IOCtl(FIONBIO,&dwArgument);

	// Receive the packet
	rhs.SerializePacket(this,false);

	// Disable blocking mode and reinstate events
	dwArgument = FD_CONNECT | FD_READ | FD_CLOSE;
	IOCtl(FIONBIO,&dwArgument);
	AsyncSelect(dwArgument);

	// Return a reference to "this" class for chaining
	return *this;
}

// Send the buffer of data
int CTCPSocket::Send(const void * lpBuf,int nBufLen,int nFlags)
{
	// Test for a function pointer to the TransmitPackets function
	if (!TransmitPackets)
		return 0;

	// Create the transmit buffer data structure
	_TRANSMIT_PACKETS_ELEMENT TPE;
	TPE.dwElFlags = TP_ELEMENT_MEMORY;
	TPE.cLength = nBufLen;
	TPE.pBuffer = (PVOID)lpBuf;

	// Transmit the packet
	BOOL bTransmit = TransmitPackets(m_hSocket,&TPE,1,nBufLen,NULL,TF_USE_KERNEL_APC);
	return nBufLen;
}

// Receive the buffer of data
int CTCPSocket::Receive(void * lpBuf,int nBufLen,int nFlags)
{
	int nTotalRecv = 0,nRecvAttempt = 0;
	while (nTotalRecv < nBufLen)
	{
		int nRecv = baseclass::Receive((LPSTR)lpBuf + nTotalRecv,nBufLen - nTotalRecv,nFlags);
		if (nRecv == SOCKET_ERROR)
			return nRecv;
		nTotalRecv += nRecv;

		// Test for a dropped connection during a receive
		if (nRecv == 0)
		{
			nRecvAttempt++;
			Sleep(300);
		}
		if (nRecvAttempt == 10)
			return 0;
	}
	return nTotalRecv;
}